#include <QFile>

#include "gsl/gsl_statistics.h"
#include "algorithm"
#include "network.h"
#include "plotter.h"
#include <QRegExp>

bool compareNodes(Node * a, Node * b)
{
    return a->stats->median < b->stats->median;
}

Network::Network()
{
     nodes = new QVector<Node *>();
     samples = new QVector<QVector<double> *>();
}

Network::~Network()
{
    for(int i = 0; i < nodes->count(); i++)
    {
        delete nodes->at(i);
    }
    delete nodes;

    for(int i = 0; i < samples->count(); i++)
    {
        delete samples->at(i);
    }
    delete samples;
}


QVector<Node *> * Network::getNodes()
{
    return nodes;
}

int Network::numNodes()
{
    return nodes->count();
}

Node * Network::getNode(int i)
{
    if(i < 0 || i >= nodes->count())
        return NULL;

    return nodes->at(i);
}


void Network::addNode(Node * n)
{
    nodes->append(n);
}


void Network::addSubnet(Network * n)
{
    QVector<Node *> *ns = n->getNodes();
    foreach(Node * node, *ns)
    {
        nodes->append(node);
    }
}


void Network::generateSamples(int N, int burn)
{
    // allocate vectors appropriately
    int count = nodes->count();

    if(count < samples->count())
    {
        samples->remove(count, samples->count() - count);
    }

    for(int i = 0; i < count; i++)
    {
        if(i < samples->count())
        {
            samples->at(i)->clear();
        }
        else
        {
            samples->append(new QVector<double>());
        }
    }

    // do burn and sample
    int total = N + burn;

    // do sampling (Gibbs)
    for(int i = 0; i < total; i++)
    {
        for(int j = 0; j < count; j++)
        {
            // can ignore those nodes that are constant / have no distribution
            if(nodes->at(j)->getIsObservation())
            {
                continue;
            }

            // always generate a value (since values propogate throughout the network)

            double value = nodes->at(j)->generateSample();

            // only store non-burn values
            if(i >= burn)
            {
                samples->at(j)->append(value);
            }
        }
    }
}



QVector<double> * Network::getSamplesForNode(int i)
{
    if(i < 0 || i >= samples->count())
        return NULL;

    return samples->at(i);
}


void Network::computeStats(int numBins, bool saveStats)
{
    Node * node;
    double binSize;
    QVector<int> * histogram;
    QVector<double> * histogramValue;
    QVector<double> * sample;        

    // sort data
    for(int j = 0; j < nodes->count(); j++)
    {
        node = nodes->at(j);
        if(node->getIsObservation())
            continue;

        sample = samples->at(j);
        qSort(*sample);

        // compute point estimators
        node->stats->cnt          = numBins;
        node->stats->min          = sample->at(0);
        node->stats->fifth        = sample->at((int)(0.05 * sample->count()));
        node->stats->ninetyFifth  = sample->at((int)(0.95 * sample->count()));
        node->stats->max          = sample->at(sample->count() - 1);
        node->stats->mean         = gsl_stats_mean(sample->data(), 1, sample->count());
        node->stats->median       = gsl_stats_median_from_sorted_data(sample->data(), 1, sample->count());
        node->stats->variance     = gsl_stats_variance(sample->data(), 1, sample->count());

        // init histogram
        binSize = (node->stats->max - node->stats->min) / (double)numBins;
        histogram = node->stats->histogram;
        histogramValue = node->stats->histogramValue;

        // compute histogram        
        double value;
        int    curBin = 0;
        double curVal = node->stats->min;
        histogram->append(0);
        histogramValue->append(curVal);
        for(int i = 0; i < sample->count(); i++)
        {
            value = sample->at(i);
            while(value > (curVal + binSize) && curBin < numBins)
            {
                curVal += binSize;
                histogram->append(0);
                histogramValue->append(curVal);
                curBin++;
            }
            (*histogram)[curBin]++;
        }       

        // save stats to file
        if(saveStats)
        {
            node->stats->saveStats(QString("out/") + node->getLabel() + QString("_stats.txt"));
        }
    }

    std::sort( nodes->begin(), nodes->end(), compareNodes );
}


void Network::print()
{
    QVector<double> * sample;
    int numSamples;
    for(int i = 0; i < nodes->count(); i++)
    {
        if(nodes->at(i)->getIsObservation())
            continue;

        sample     = samples->at(i);
        numSamples = sample->count();

        qDebug() << i << nodes->at(i)->getLabel() << ": " << nodes->at(i)->stats->median << ";  90% interval: (" << nodes->at(i)->stats->fifth << nodes->at(i)->stats->ninetyFifth << ")";
        //QVector<Node *> * parents = nodes->at(i)->getParentNodes();
       /* qDebug() << "\t" << nodes->at(i)->getValue();
        for(int j = 0; j < parents->count(); j++)
        {
            Node * parent = parents->at(j);
            qDebug() << "\t" << parent->getLabel() << parent->getValue();
        }*/
    }
}


void Network::saveSamples(QString path)
{
    QFile file(path);
    file.open(QIODevice::WriteOnly | QIODevice::Text);
    QTextStream out(&file);

    // header row
    for(int j = 0; j < nodes->count(); j++)
    {
        out << nodes->at(j)->getLabel() << " ";
    }
    out << endl;

    // data
    for(int i = 0; i < samples->at(0)->count(); i++)
    {
        for(int k = 0; k < nodes->count(); k++)
        {
           if(nodes->at(k)->getIsObservation())
               continue;

           out << samples->at(k)->at(i) << " ";
        }
        out << endl;
    }

    file.close();
}

void Network::plotMixingGraphs(QString path)
{
    for(int i = 0; i < nodes->count(); i++)
    {
        if(nodes->at(i)->getIsObservation())
            continue;
        QString label = nodes->at(i)->getLabel();
        QString tmp   = path;
        QString out   = tmp.append("/").append(label).append("mixing");
        Plotter::plotY(getSamplesForNode(i),
                       Plotter::LINES,
                       label,
                       out);
    }
}

void Network::plotHistograms(QString path)
{
    for(int i = 0; i < nodes->count(); i++)
    {
        if(nodes->at(i)->getIsObservation())
            continue;
        QString label = nodes->at(i)->getLabel();
        QString tmp   = path;
        QString out   = tmp.append("/").append(label).append("histogram");
        Plotter::plotXY(nodes->at(i)->stats->histogramValue,
                        nodes->at(i)->stats->histogram,
                        Plotter::IMPULSES,
                        label,
                        out);
    }
}



/**
 * Prunes a bayesian network according to the d-separability algorithm and a slight modification
 * to account for implicit evidence (all nodes have implicit evidence).
 */

QVector<Node *> * Network::prune(QVector<Node *> * queryNodes, QVector<Node *> * observedNodes)
{
    Node * node;

    QVector<Node *> observedAndAncestors;
    QVector<Node *> todo(*observedNodes);


    foreach(Node * node, *nodes)
    {
        node->isNecessary = false;
    }


// =============== Phase 1: Put all parents of observations in ancestor list ===============

    while(!todo.isEmpty())
    {
        node = todo.last();
        todo.pop_back();

        if(!observedAndAncestors.contains(node))
        {
            todo += node->getExplicitParentNodes();
        }

        observedAndAncestors.push_back(node);
    }


// =============== Phase 2: Add nodes along active trails ==================================

    QVector<Node *> * necessary = new QVector<Node *>();
    QVector<NodeDirectionTuple> visited;
    QVector<NodeDirectionTuple> L;
    Direction dir;
    bool ancestorOfEvidence, notObserved;


    // Initialize visit list with query nodes since that is where the traversal begins
    foreach(Node * node, *queryNodes)
    {
        L.push_back(NodeDirectionTuple(node));
    }

    // as long as there are nodes to visit
    while(L.count())
    {
        NodeDirectionTuple currentNode = L.last();
        L.pop_back();

        node        = currentNode.node;
        dir         = currentNode.dir;

        if(visited.contains(currentNode))
            continue;

        visited.push_back(currentNode);

        ancestorOfEvidence = observedAndAncestors.contains(node);
        notObserved        = !observedNodes->contains(node);

        // Current node is only nesccessary if it is an ancestor of evidence (going down)
        // OR it is a decendant of evidence (going up, which is always due to implicit evidence)
        if((dir == DOWN && ancestorOfEvidence
             || dir == UP)
            && !necessary->contains(node))
        {
            necessary->append(node);
            node->isNecessary = true;
        }

        // Came to this node from a child node (can only continue to parents or children if not observed)
        if(dir == UP && notObserved)
        {
            // add parent nodes to todo list
            Network::addNodeTuplesToList(&L, node->getExplicitParentNodes(), UP);

            // add child nodes to todo list
            Network::addNodeTuplesToList(&L, node->getChildNodes(), DOWN);
        }

        // Came to this node from a parent node
        else if(dir == DOWN)
        {
            if(ancestorOfEvidence)
                Network::addNodeTuplesToList(&L, node->getExplicitParentNodes(), UP);

            if(notObserved)
                Network::addNodeTuplesToList(&L, node->getChildNodes(), DOWN);
        }
    }

    return necessary;
}


void Network::addNodeTuplesToList(QVector<NodeDirectionTuple> * list, QVector<Node*> nodes, Direction dir)
{
    foreach(Node * node, nodes)
    {
        list->push_back(NodeDirectionTuple(node, dir));
    }
}

void Network::addNodeTuplesToList(QVector<NodeDirectionTuple> * list, QVector<Node*> * nodes, Direction dir)
{
    foreach(Node * node, *nodes)
    {
        list->push_back(NodeDirectionTuple(node, dir));
    }
}

void Network::toDotFile(QString filename, QVector<Node *> * highlightNodes, QVector<Node *> * keepNodes)
{
   QFile outFile;
   outFile.setFileName(filename);
   outFile.open(QIODevice::WriteOnly);
   QTextStream out(&outFile);

   out << "digraph unix {\n";
   out << "size=\"200,200\";\n";

   foreach (Node* node, *nodes)
   {
        if(node->getIsObservation())
        {
            out << node->getLabel() << " [color=red, style=filled];\n";
        }
        else if (highlightNodes != NULL && highlightNodes->contains(node))
        {
            out << node->getLabel() << " [color=green, style=filled];\n";
        }
        else if (keepNodes == NULL || keepNodes->contains(node))
        {
            out << node->getLabel() << " [color=black];\n";
        }
   }

   foreach (Node* node, *nodes)
   {
       if(!node->isNecessary)
           continue;

       QString name = node->getLabel();
       foreach (Node* child, *(node->getChildNodes()))
       {
           if(child->isNecessary)
           {
              out << "\"" << name << "\"" << " -> " << "\"" << child->getLabel() << "\"" << ";\n";
           }
       }
   }

   out << "}";
   outFile.close();
   QString newFile = filename;
   newFile.replace(".out", ".png");
   QString command = "dot -Tpng " + filename + " -o " + newFile;
   system(command.toAscii().data());
   command = "rm " + filename;
   system(command.toAscii().data());
}
